/*
 * Copyright 2024 ywcai
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

`include "defines.v"
`timescale 1ns/1ps

module mem(
	input	wire[`InstDataBus]	inst_i,
	input	wire[`MemAddrBus]	pc_i,
	input	wire				regwrite_en_i,
	input	wire[`RegAddrBus]	regwrite_addr_i,
	input	wire[`RegDataBus]	mem_read_data_i,
	input	wire[`RegDataBus]	rs1_data_i,
	input	wire[`RegDataBus]	rs2_data_i,
	input	wire[`RegDataBus]	rs1_imm_value_i,
	input	wire[`RegDataBus]	ex_result_i,
	input	wire[2:0]			mux_result_src_i,
	input	wire				mux_jump_i,
	input	wire				mux_loadtype_i,
	input	wire				mux_storetype_i,
	input	wire[`MemAddrBus]	pc_target_i,
	input	wire				csr_mie_i,
	input	wire				csr_wr_en_i,
	input	wire[`MemAddrBus]	csr_wr_addr_i,
	input	wire[`RegDataBus]	csr_wr_data_i,
	input	wire[`RegDataBus]	csr_rd_data_i,
    input	wire[`RegDataBus]	csr_mtvec_i,
    input	wire[`RegDataBus]	csr_mepc_i,
    input	wire[`RegDataBus]	csr_mstatus_i,
	input	wire[`IrqBus]		irq_i,
	input	wire[`TrapBus]		trap_code_i,
	input	wire				trap_assert_i,
	input	wire[`MemAddrBus]	trap_entry_i,

	output	wire[`MemAddrBus]	pc_o,
	output	wire[`MemAddrBus]	trap_pc_o,
	output	wire				regwrite_en_o,
	output	wire[`RegAddrBus]	regwrite_addr_o,
	output	wire[`RegDataBus]	ex_result_o,
	output	wire[`MemAddrBus]	pc_target_o,
	output	wire[2:0]			mux_result_src_o,
	output	wire[`RegDataBus]	mem_read_data_o,
	output	wire				csr_mie_o,
	output	wire				csr_wr_en_o,
	output	wire[`MemAddrBus]	csr_wr_addr_o,
	output	wire[`RegDataBus]	csr_wr_data_o,
	output	wire[`MemAddrBus]	csr_rd_addr_o,
	output	wire[`RegDataBus]	csr_rd_data_o,
    output	wire[`RegDataBus]	csr_mtvec_o,
    output	wire[`RegDataBus]	csr_mepc_o,
    output	wire[`RegDataBus]	csr_mstatus_o,
	output	wire[`TrapBus]		trap_code_o,
	output	wire[`MemAddrBus]	trap_entry_o
    );

    wire[2:0] funct3	= inst_i[14:12];
    wire[6:0] opcode	= inst_i[6:0];

	wire i_inst_load	= (opcode == `INST_LOAD);

	assign pc_o			= pc_i;
	assign trap_pc_o	= pc_i;
	assign trap_entry_o	= trap_entry_i;

	assign regwrite_en_o	= trap_assert_i ? `DISABLE : regwrite_en_i;
	assign regwrite_addr_o	= regwrite_addr_i;
	assign ex_result_o		= ex_result_i;
	assign mux_result_src_o	= mux_result_src_i;
	assign pc_target_o		= pc_target_i;

	wire[`MemAddrBus] load_addr	= ex_result_i;

	wire i_inst_lb	= (i_inst_load & (funct3 == `FUNCT3_000));
	wire i_inst_lbu	= (i_inst_load & (funct3 == `FUNCT3_100));
	wire i_inst_lh	= (i_inst_load & (funct3 == `FUNCT3_001));
	wire i_inst_lhu	= (i_inst_load & (funct3 == `FUNCT3_101));
	wire i_inst_lw	= (i_inst_load & (funct3 == `FUNCT3_010));
	wire i_inst_lwu	= (i_inst_load & (funct3 == `FUNCT3_110));
	wire i_inst_ld	= (i_inst_load & (funct3 == `FUNCT3_011));

	wire[2:0] byte_idx = load_addr[2:0] & 3'b111;
	wire[`RegDataBus] lb_byte1_ext = {{56{mem_read_data_i[7]}}, mem_read_data_i[7:0]};
	wire[`RegDataBus] lb_byte2_ext = {{56{mem_read_data_i[15]}}, mem_read_data_i[15:8]};
	wire[`RegDataBus] lb_byte3_ext = {{56{mem_read_data_i[23]}}, mem_read_data_i[23:16]};
	wire[`RegDataBus] lb_byte4_ext = {{56{mem_read_data_i[31]}}, mem_read_data_i[31:24]};
	wire[`RegDataBus] lb_byte5_ext = {{56{mem_read_data_i[39]}}, mem_read_data_i[39:32]};
	wire[`RegDataBus] lb_byte6_ext = {{56{mem_read_data_i[47]}}, mem_read_data_i[47:40]};
	wire[`RegDataBus] lb_byte7_ext = {{56{mem_read_data_i[55]}}, mem_read_data_i[55:48]};
	wire[`RegDataBus] lb_byte8_ext = {{56{mem_read_data_i[63]}}, mem_read_data_i[63:56]};
	wire[`RegDataBus] data_lb =
        ({`XLEN{byte_idx == 3'b000}} & lb_byte1_ext)
		| ({`XLEN{byte_idx == 3'b001}} & lb_byte2_ext)
		| ({`XLEN{byte_idx == 3'b010}} & lb_byte3_ext)
		| ({`XLEN{byte_idx == 3'b011}} & lb_byte4_ext)
		| ({`XLEN{byte_idx == 3'b100}} & lb_byte5_ext)
		| ({`XLEN{byte_idx == 3'b101}} & lb_byte6_ext)
		| ({`XLEN{byte_idx == 3'b110}} & lb_byte7_ext)
		| ({`XLEN{byte_idx == 3'b111}} & lb_byte8_ext);
	wire[`RegDataBus] lbu_byte1_ext = {{56{1'b0}}, mem_read_data_i[7:0]};
	wire[`RegDataBus] lbu_byte2_ext = {{56{1'b0}}, mem_read_data_i[15:8]};
	wire[`RegDataBus] lbu_byte3_ext = {{56{1'b0}}, mem_read_data_i[23:16]};
	wire[`RegDataBus] lbu_byte4_ext = {{56{1'b0}}, mem_read_data_i[31:24]};
	wire[`RegDataBus] lbu_byte5_ext = {{56{1'b0}}, mem_read_data_i[39:32]};
	wire[`RegDataBus] lbu_byte6_ext = {{56{1'b0}}, mem_read_data_i[47:40]};
	wire[`RegDataBus] lbu_byte7_ext = {{56{1'b0}}, mem_read_data_i[55:48]};
	wire[`RegDataBus] lbu_byte8_ext = {{56{1'b0}}, mem_read_data_i[63:56]};
	wire[`RegDataBus] data_lbu =
        ({`XLEN{byte_idx == 3'b000}} & lbu_byte1_ext)
		| ({`XLEN{byte_idx == 3'b001}} & lbu_byte2_ext)
		| ({`XLEN{byte_idx == 3'b010}} & lbu_byte3_ext)
		| ({`XLEN{byte_idx == 3'b011}} & lbu_byte4_ext)
		| ({`XLEN{byte_idx == 3'b100}} & lbu_byte5_ext)
		| ({`XLEN{byte_idx == 3'b101}} & lbu_byte6_ext)
		| ({`XLEN{byte_idx == 3'b110}} & lbu_byte7_ext)
		| ({`XLEN{byte_idx == 3'b111}} & lbu_byte8_ext);

	wire[1:0] half_idx = load_addr[2:1] & 2'b11;
	wire[`RegDataBus] lh_word1_ext = {{48{mem_read_data_i[15]}}, mem_read_data_i[15:0]};
	wire[`RegDataBus] lh_word2_ext = {{48{mem_read_data_i[31]}}, mem_read_data_i[31:16]};
	wire[`RegDataBus] lh_word3_ext = {{48{mem_read_data_i[47]}}, mem_read_data_i[47:32]};
	wire[`RegDataBus] lh_word4_ext = {{48{mem_read_data_i[63]}}, mem_read_data_i[63:48]};
	wire[`RegDataBus] data_lh =
        ({`XLEN{half_idx == 2'b00}} & lh_word1_ext)
		| ({`XLEN{half_idx == 2'b01}} & lh_word2_ext)
		| ({`XLEN{half_idx == 2'b10}} & lh_word3_ext)
		| ({`XLEN{half_idx == 2'b11}} & lh_word4_ext);
	wire[`RegDataBus] lhu_word1_ext = {{48{1'b0}}, mem_read_data_i[15:0]};
	wire[`RegDataBus] lhu_word2_ext = {{48{1'b0}}, mem_read_data_i[31:16]};
	wire[`RegDataBus] lhu_word3_ext = {{48{1'b0}}, mem_read_data_i[47:32]};
	wire[`RegDataBus] lhu_word4_ext = {{48{1'b0}}, mem_read_data_i[63:48]};
	wire[`RegDataBus] data_lhu =
        ({`XLEN{half_idx == 2'b00}} & lhu_word1_ext)
		| ({`XLEN{half_idx == 2'b01}} & lhu_word2_ext)
		| ({`XLEN{half_idx == 2'b10}} & lhu_word3_ext)
		| ({`XLEN{half_idx == 2'b11}} & lhu_word4_ext);

	wire word_idx = load_addr[2] & 1'b1;
	wire[`RegDataBus] lw_word1_ext = {{32{mem_read_data_i[31]}}, mem_read_data_i[31:0]};
	wire[`RegDataBus] lw_word2_ext = {{32{mem_read_data_i[63]}}, mem_read_data_i[63:32]};
	wire[`RegDataBus] data_lw =
        ({`XLEN{word_idx == 1'b0}} & lw_word1_ext)
		| ({`XLEN{word_idx == 1'b1}} & lw_word2_ext);
	wire[`RegDataBus] lwu_word1_ext = {{32{1'b0}}, mem_read_data_i[31:0]};
	wire[`RegDataBus] lwu_word2_ext = {{32{1'b0}}, mem_read_data_i[63:32]};
	wire[`RegDataBus] data_lwu =
        ({`XLEN{word_idx == 1'b0}} & lwu_word1_ext)
		| ({`XLEN{word_idx == 1'b1}} & lwu_word2_ext);

	wire[`RegDataBus] data_ld = mem_read_data_i;

	assign mem_read_data_o = ({`XLEN{i_inst_lb}} & data_lb)
		| ({`XLEN{i_inst_lbu}} & data_lbu)
		| ({`XLEN{i_inst_lh}} & data_lh)
		| ({`XLEN{i_inst_lhu}} & data_lhu)
		| ({`XLEN{i_inst_lw}} & data_lw)
		| ({`XLEN{i_inst_lwu}} & data_lwu)
		| ({`XLEN{i_inst_ld}} & data_ld);

	wire i_inst_csrrw  = (funct3 == `FUNCT3_001);
	wire i_inst_csrrs  = (funct3 == `FUNCT3_010);
	wire i_inst_csrrc  = (funct3 == `FUNCT3_011);
	wire i_inst_csrrwi = (funct3 == `FUNCT3_101);
	wire i_inst_csrrsi = (funct3 == `FUNCT3_110);
	wire i_inst_csrrci = (funct3 == `FUNCT3_111);

	wire csr_inst_wr_en		= (i_inst_csrrw & i_inst_csrrwi);
	wire[`MemAddrBus] csr_inst_addr = {28'h0, inst_i[31:20]};
	assign csr_mie_o		= csr_mie_i;
	assign csr_wr_en_o		= trap_assert_i ? csr_wr_en_i : csr_inst_wr_en;
	assign csr_wr_addr_o	= trap_assert_i ? csr_wr_addr_i : csr_inst_addr;
	assign csr_rd_data_o	= csr_rd_data_i;
	assign csr_mtvec_o		= csr_mtvec_i;
	assign csr_mepc_o		= csr_mepc_i;
	assign csr_mstatus_o	= csr_mstatus_i;
	assign csr_rd_addr_o	= csr_inst_addr;

	wire[`RegDataBus] csrrs_data	= rs1_data_i | csr_rd_data_i;
	wire[`RegDataBus] csrrc_data	= rs1_data_i & (~csr_rd_data_i);
	wire[`RegDataBus] csrrsi_data	= rs1_imm_value_i | csr_rd_data_i;
	wire[`RegDataBus] csrrci_data	= rs1_imm_value_i & (~csr_rd_data_i);

	wire[`RegDataBus] csr_inst_wr_data = ({`XLEN{i_inst_csrrw}} & rs1_data_i)
		| ({`XLEN{i_inst_csrrs}} & csrrs_data)
		| ({`XLEN{i_inst_csrrc}} & csrrc_data)
		| ({`XLEN{i_inst_csrrwi}} & rs1_imm_value_i)
		| ({`XLEN{i_inst_csrrsi}} & csrrsi_data)
		| ({`XLEN{i_inst_csrrci}} & csrrci_data);

	assign csr_wr_data_o = trap_assert_i ? csr_wr_data_i : csr_inst_wr_data;

    assign trap_code_o = trap_code_i;

endmodule
